import asyncio

from fastapi import APIRouter, Depends, HTTPException
from neo4j import AsyncDriver
from starlette.responses import StreamingResponse
from ..services.rag_service import RAGService
from ..services.graph_query_service import GraphQueryService
from ..schemas.schemas import RagQueryRequest
from sqlalchemy.ext.asyncio import AsyncSession
from core.dependencies import get_db,get_neo4j_driver
from ...knowbase.services.neo4j_service import Neo4jService
import json

router = APIRouter(prefix="/chat",tags=["chat"])

#依赖项：获取Neo4j服务实例

async def get_neo4j_service(driver:AsyncDriver = Depends(get_neo4j_driver)):
    return Neo4jService(driver)

rag_service = RAGService()
graph_query_service = GraphQueryService()


@router.post("/rag_query_stream")
async def rag_query_stream(
        request: RagQueryRequest,
        session: AsyncSession = Depends(get_db)
):
    """RAG查询流式端点"""

    async def event_generator():
        """生成SSE事件流"""
        async for chunk in rag_service.process_rag_query(
                know_base_ids=request.knowledge_bases,
                question=request.question,
                session=session
        ):
            yield chunk

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )

@router.post("/graph_query")
async def graph_query(question: str, neo4j_service: Neo4jService = Depends(get_neo4j_service)):
    keyword =  await graph_query_service.extract_keywords(question)
    try:
        results = await neo4j_service.keyword_search(keyword)
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/unified_rag_query")
async def unifed_rag_query(
        request: RagQueryRequest,
        session:AsyncSession = Depends(get_db),
        neo4j_service: Neo4jService = Depends(get_neo4j_service)
    ):
        # 统一的RAG查询接口，同时返回流式文本响应和图谱数据的SSE事件流
        async def event_generator():
            """生成SSE事件流 - 修正版本"""
            try:
                # 直接使用 async for 迭代异步生成器
                async for chunk in rag_service.process_rag_query(
                        know_base_ids=request.knowledge_bases,
                        question=request.question,
                        session=session
                ):

                    yield chunk

                # 然后处理知识图谱数据
                keyword = await graph_query_service.extract_keywords(request.question)
                graph_results = await neo4j_service.keyword_search(keyword)

                # 发送图谱数据事件
                graph_event = {
                    "type": "graph_data",
                    "graph_data": graph_results
                }
                yield f"data: {json.dumps(graph_event,ensure_ascii=False)}\n\n"

            except Exception as e:
                error_event = {
                    "event": "error",
                    "": str(e)
                }
                yield f"data: {json.dumps(error_event)}\n\n"

        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no"  # 禁用Nginx缓冲
            }
        )

